#ifndef BOUDNARY_TOOLS_HPP_
#define BOUDNARY_TOOLS_HPP_

#include <Eigen/Core>
#include <Eigen/Sparse>
#include <Eigen/src/Core/Matrix.h>
#include <array>
#include <vector>
#include "Function.hpp"

namespace TetTools {

inline std::vector<int> get_boundary(const Eigen::MatrixXi &T) {
  int num_of_cells = T.rows();
  // 遍历四面体单元查找边界
  std::vector<std::array<int, 3>>
      boundary; // 记录所有网格单元面的顺序节点三元组
  std::vector<int> boundary_nodes; //  按顺序记录所有网格边界节点
  for (int n = 0; n < num_of_cells; n++) {
    std::vector<int> v = {T.row(n)[0], T.row(n)[1], T.row(n)[2], T.row(n)[3]};
    std::sort(v.begin(), v.end());
    boundary.push_back({v[0], v[1], v[2]});
    boundary.push_back({v[0], v[1], v[3]});
    boundary.push_back({v[0], v[2], v[3]});
    boundary.push_back({v[1], v[2], v[3]});
  }
  std::sort(boundary.begin(), boundary.end());
  for (int n = 0; n < boundary.size() - 1; n++) {
    if (boundary[n] == boundary[n + 1]) {
      n++;
    } else {
      for (int i = 0; i < 3; i++) {
        boundary_nodes.push_back(boundary[n][i]);
      }
    }
  }
  if (boundary[boundary.size() - 1] != boundary[boundary.size() - 2]) {
    for (int i = 0; i < 3; i++) {
      boundary_nodes.push_back(boundary[boundary.size()-1][i]); // /? might be wrong
    }
  }
  std::sort(boundary_nodes.begin(), boundary_nodes.end());
  boundary_nodes.erase(
      std::unique(boundary_nodes.begin(), boundary_nodes.end()),
      boundary_nodes.end());
  return boundary_nodes;
}

inline void set_boundary(const Function<3> &bd_ftn, const std::vector<int> &boundary_id, const Eigen::MatrixXd &V,
                         Eigen::SparseMatrix<double> &system_matrix,
                         Eigen::VectorXd &system_rhs) {
  int num_of_nodes = system_rhs.rows();
  std::vector<int> is_boundary(num_of_nodes, 0);
  for (int i : boundary_id) {
    is_boundary[i] = 1;
  }
  /* 右端项边界置为边界值，非边界减去边界值 */
  for (int i : boundary_id) {
    Eigen::Vector<double, 3> point = V.row(i); 
    for (int j=0; j<num_of_nodes; j++) {
      system_rhs(j) -= system_matrix.coeff(j, i)*bd_ftn(point); 
    }
  }
  for (int i : boundary_id) {
    system_rhs(i) = bd_ftn(V.row(i));
  }
  /* 矩阵在边界节点的行和列归0 */
  system_matrix.prune([&is_boundary](int i, int j, double) {
    return (is_boundary[i] == 0 && is_boundary[j] == 0);
  });
  /* 矩阵边界节点对角元置为1 */
  std::vector<Eigen::Triplet<double>> triplets_list;
  for (int i : boundary_id) {
    triplets_list.push_back(Eigen::Triplet<double>(i,i,1));
  }
  Eigen::SparseMatrix<double> M_tmp(num_of_nodes,num_of_nodes);
  M_tmp.setZero();
  M_tmp.setFromTriplets(triplets_list.begin(), triplets_list.end());
  system_matrix = system_matrix + M_tmp;
}
} // namespace TetTools


namespace HexTools {
inline std::vector<int> get_boundary(const Eigen::MatrixXi &T) {
  int num_of_cells = T.rows();
  // 遍历六面体单元查找边界
  std::vector<std::array<int, 4>>
      boundary; // 记录所有网格单元面的顺序节点四元组
  std::vector<int> boundary_nodes; //  按顺序记录所有网格边界节点
  for (int n = 0; n < num_of_cells; n++) {
    std::array<int, 4> v;
    v = {T.row(n)[0], T.row(n)[1], T.row(n)[2], T.row(n)[3]};
    std::sort(v.begin(), v.end());
    boundary.push_back(v);
    v = {T.row(n)[0], T.row(n)[1], T.row(n)[4], T.row(n)[5]};
    std::sort(v.begin(), v.end());
    boundary.push_back(v);
    v = {T.row(n)[0], T.row(n)[3], T.row(n)[4], T.row(n)[7]};
    std::sort(v.begin(), v.end());
    boundary.push_back(v);
    v = {T.row(n)[1], T.row(n)[2], T.row(n)[5], T.row(n)[6]};
    std::sort(v.begin(), v.end());
    boundary.push_back(v);
    v = {T.row(n)[2], T.row(n)[3], T.row(n)[6], T.row(n)[7]};
    std::sort(v.begin(), v.end());
    boundary.push_back(v);
    v = {T.row(n)[4], T.row(n)[5], T.row(n)[6], T.row(n)[7]};
    std::sort(v.begin(), v.end());
    boundary.push_back(v);
  }
  std::sort(boundary.begin(), boundary.end());
  for (int n = 0; n < boundary.size() - 1; n++) {
    if (boundary[n] == boundary[n + 1]) {
      n++;
    } else {
      for (int i = 0; i < 4; i++) {
        boundary_nodes.push_back(boundary[n][i]);
      }
    }
  }
  if (boundary[boundary.size() - 1] != boundary[boundary.size() - 2]) {
    for (int i = 0; i < 4; i++) {
      boundary_nodes.push_back(boundary[boundary.size()-1][i]);
    }
  }
  std::sort(boundary_nodes.begin(), boundary_nodes.end());
  boundary_nodes.erase(
      std::unique(boundary_nodes.begin(), boundary_nodes.end()),
      boundary_nodes.end());
  return boundary_nodes;
}

inline void set_boundary(const Function<3> &bd_ftn, const std::vector<int> &boundary_id, const Eigen::MatrixXd &V,
                         Eigen::SparseMatrix<double> &system_matrix,
                         Eigen::VectorXd &system_rhs) {
  int num_of_nodes = system_rhs.rows();
  std::vector<int> is_boundary(num_of_nodes, 0);
  for (int i : boundary_id) {
    is_boundary[i] = 1;
  }
  /* 右端项边界置为边界值，非边界减去边界值 */
  for (int i : boundary_id) {
    Eigen::Vector<double, 3> point = V.row(i); 
    for (int j=0; j<num_of_nodes; j++) {
      system_rhs(j) -= system_matrix.coeff(j, i)*bd_ftn(point); 
    }
  }
  for (int i : boundary_id) {
    system_rhs(i) = bd_ftn(V.row(i));
  }
  /* 矩阵在边界节点的行和列归0 */
  system_matrix.prune([&is_boundary](int i, int j, double) {
    return (is_boundary[i] == 0 && is_boundary[j] == 0);
  });
  /* 矩阵边界节点对角元置为1 */
  std::vector<Eigen::Triplet<double>> triplets_list;
  for (int i : boundary_id) {
    triplets_list.push_back(Eigen::Triplet<double>(i,i,1));
  }
  Eigen::SparseMatrix<double> M_tmp(num_of_nodes,num_of_nodes);
  M_tmp.setZero();
  M_tmp.setFromTriplets(triplets_list.begin(), triplets_list.end());
  system_matrix = system_matrix + M_tmp;
}
} // namespace HexTools
 
#endif